#ifndef FLASHROS_INFERENCE_OUTPUT_PIPELINE_H
#define FLASHROS_INFERENCE_OUTPUT_PIPELINE_H

#include <cv_bridge/cv_bridge.h>
#include <image_transport/image_transport.h>
#include <ros/ros.h>
#include <spdlog/spdlog.h>

#include <chrono>
#include <mutex>
#include <opencv2/opencv.hpp>
#include <string>
#include <thread>

#include "flashai/commons/box.h"
#include "flashai/commons/visualize.h"
#include "flashros_commons/BoxesMsg.h"

namespace flashros {
namespace inference {

class DetectionOutputInterface {
 public:
  virtual void Write(cv::Mat &frame,
                     const std::vector<flashai::commons::Box> &boxes) = 0;
  virtual void Release() = 0;
};

class VisualizeDetectionOutput : public DetectionOutputInterface {
 public:
  VisualizeDetectionOutput(INIReader &reader, std::string section,
                           ros::NodeHandle &nh, std::string data_root) {
    Init(reader, section, nh, data_root);
  }

  void Init(INIReader &reader, std::string section, ros::NodeHandle &nh,
            std::string data_root) {
    ROS_WARN_STREAM("start init visualizer with section "
                    << section << " and data root " << data_root);
    visualize_.Init(reader, section, data_root);
  }

  void Write(cv::Mat &frame, const std::vector<flashai::commons::Box> &boxes) {
    visualize_.Draw(frame, boxes);
  }

  void Release() {}

 private:
  flashai::utils::DetectionVisualizer visualize_;
};

class RosDetectionOutput : public DetectionOutputInterface {
 public:
  RosDetectionOutput() {}
  RosDetectionOutput(std::string detection_topic_name, ros::NodeHandle &nh) {
    Init(detection_topic_name, nh);
  }
  void Init(std::string detection_topic_name, ros::NodeHandle &nh) {
    SPDLOG_INFO("Start init ros output with topic name '{}'",
                detection_topic_name);
    publisher_ =
        nh.advertise<flashros_commons::BoxesMsg>(detection_topic_name, 1);
  }

  void Write(cv::Mat &frame,
             const std::vector<flashai::commons::Box> &boxes) override {
    flashros_commons::BoxesMsg output;

    // init head
    output.header.stamp = ros::Time::now();

    // init bboxes
    output.num_objects = boxes.size();
    for (auto it = boxes.begin(); it < boxes.end(); it++) {
      output.bboxes.push_back((*it).x1);
      output.bboxes.push_back((*it).y1);
      output.bboxes.push_back((*it).x2);
      output.bboxes.push_back((*it).y2);
      output.socres.push_back((*it).score);
      output.classes.push_back((*it).class_id);
      output.distances.push_back((*it).distance);
      output.texts.push_back((*it).text);
    }

    SPDLOG_DEBUG("Ros Output Publish");
    publisher_.publish(output);
    SPDLOG_DEBUG("Ros Output Publish {} boxes", boxes.size());
  }

  void Release() override {
    publisher_.shutdown();
    SPDLOG_DEBUG("Ros output release");
  }

 private:
  ros::Publisher publisher_;
};

}  // namespace inference
}  // namespace flashros
#endif  // FLASHROS_INFERENCE_OUTPUT_PIPELINE_H